import pymysql


# 公共的断言的方法：响应状态码、success、code、message
def common_assert(test_case, response, status_code, success, code, message):
    # 响应状态码
    test_case.assertEqual(status_code, response.status_code)
    test_case.assertEqual(success, response.json().get("success"))
    test_case.assertEqual(code, response.json().get("code"))
    test_case.assertIn(message, response.json().get("message"))


class DBUtil:
    _conn = None  # 数据库连接对象

    @classmethod
    def get_conn(cls):
        """获取数据库连接对象"""
        if cls._conn is None:
            cls._conn = pymysql.connect("localhost", "root", "root", "books")
        return cls._conn

    @classmethod
    def close_conn(cls):
        """关闭数据库连接"""
        if cls._conn:
            cls._conn.close()
            cls._conn = None

    @classmethod
    def get_cursor(cls):
        """获取游标对象"""
        return cls.get_conn().cursor()

    @classmethod
    def close_cursor(cls, cursor):
        """关闭游标对象"""
        if cursor:
            cursor.close()

    @classmethod
    def get_one(cls, sql):
        """查询一条记录"""
        data = None
        cursor = None
        try:
            cursor = cls.get_cursor()
            cursor.execute(sql)
            data = cursor.fetchone()
        except Exception as e:
            print("get_one error: ", e)
        finally:
            cls.close_cursor(cursor)
            cls.close_conn()
        return data
